#Load packages
library(dplyr)
library(tidyr)
library(ggplot2)
library(latex2exp)

#Load data
load("results_cs.RData")
load("results_panel.RData")

results_cs <- results_cs %>% rename(ols = ols_error)
results_cs <- results_cs %>% rename(G = modelG_error)
results_cs <- results_cs %>% rename(g = model_g_error)


med <- results_cs %>% group_by(k) %>% summarize(across(ols:g, median))
med <- med %>% pivot_longer(c(ols, G, g), names_to = "estimator", values_to = "median")
#first percentile (can't get `across` to work for sort(.)[5])
ols_fp <- results_cs %>% group_by(k) %>% summarize(ols = sort(ols)[5])
G_fp <- results_cs %>% group_by(k) %>% summarize(G = sort(G)[5])
g_fp <- results_cs %>% group_by(k) %>% summarize(g = sort(g)[5])
fp <- left_join(ols_fp, G_fp, by = "k")
fp <- left_join(fp, g_fp, by = "k")
fp <- fp %>% pivot_longer(!k, names_to = "estimator", values_to = "fp")
#ninety-ninth percentile
ols_nnp <- results_cs %>% group_by(k) %>% summarize(ols = sort(ols)[495])
G_nnp <- results_cs %>% group_by(k) %>% summarize(G = sort(G)[495])
g_nnp <- results_cs %>% group_by(k) %>% summarize(g = sort(g)[495])
nnp <- left_join(ols_nnp, G_nnp, by = "k")
nnp <- left_join(nnp, g_nnp, by = "k")
nnp <- nnp %>% pivot_longer(!k, names_to = "estimator", values_to = "nnp")
#joins
df <- left_join(med, fp, by = c("k", "estimator"))
df <- left_join(df, nnp, by = c("k", "estimator"))

#Prepping dataframe for ggplot
df_few <- df %>% filter(estimator != "ols" & k < 21)
df_few$estimator <- as.factor(df_few$estimator)
levels(df_few$estimator) <- c(g = TeX("$\\widehat{\\gamma}$"), G = TeX("$\\widehat{\\Gamma}$"))

df_20 <- df %>% filter(estimator != "ols" & k > 19)
df_20$estimator <- as.factor(df_20$estimator)
levels(df_20$estimator) <- c(g = TeX("$\\widehat{\\gamma}$"), G = TeX("$\\widehat{\\Gamma}$"))


fig1 <- ggplot(data = df_few) + geom_point(aes(x = k, y = median))
fig1 <- fig1 + geom_errorbar(aes(x = k, y = median, ymin = fp, ymax = nnp), width = 0) +
  facet_wrap(~ estimator, labeller = label_parsed)
fig1 <- fig1 + theme_bw() + theme(panel.grid.major = element_blank(), panel.grid.minor = element_blank()) +
  ylab("Error") + xlab("Regions (5 units per region)") + geom_hline(yintercept = 0)
pdf("fig1.pdf", width = 5, height = 3)
fig1
dev.off()

fig2 <- ggplot(data = df_20) + geom_point(aes(x = k, y = median))
fig2 <- fig2 + geom_errorbar(aes(x = k, y = median, ymin = fp, ymax = nnp), width = 0) +
  facet_wrap(~ estimator, labeller = label_parsed)
fig2 <- fig2 + theme_bw() + theme(panel.grid.major = element_blank(), panel.grid.minor = element_blank()) +
  ylab("Error") + xlab("Regions (5 units per region)") + geom_hline(yintercept = 0)
pdf("fig2.pdf", width = 5, height = 3)
fig2
dev.off()


# Third row of results uses panel data
results_panel <- results_panel %>% rename(fe = fe_error)
results_panel <- results_panel %>% rename(G = modelG_error)
results_panel <- results_panel %>% rename(g = model_g_error)

med <- results_panel %>% group_by(k) %>% summarize(across(fe:g, median))
med <- med %>% pivot_longer(c(fe, G, g), names_to = "estimator", values_to = "median")
#first percentile (can't get `across` to work for sort(.)[5])
fe_fp <- results_panel %>% group_by(k) %>% summarize(fe = sort(fe)[5])
G_fp <- results_panel %>% group_by(k) %>% summarize(G = sort(G)[5])
g_fp <- results_panel %>% group_by(k) %>% summarize(g = sort(g)[5])
fp <- left_join(fe_fp, G_fp, by = "k")
fp <- left_join(fp, g_fp, by = "k")
fp <- fp %>% pivot_longer(!k, names_to = "estimator", values_to = "fp")
#ninety-ninth percentile
fe_nnp <- results_panel %>% group_by(k) %>% summarize(fe = sort(fe)[495])
G_nnp <- results_panel %>% group_by(k) %>% summarize(G = sort(G)[495])
g_nnp <- results_panel %>% group_by(k) %>% summarize(g = sort(g)[495])
nnp <- left_join(fe_nnp, G_nnp, by = "k")
nnp <- left_join(nnp, g_nnp, by = "k")
nnp <- nnp %>% pivot_longer(!k, names_to = "estimator", values_to = "nnp")
#joins
df2 <- left_join(med, fp, by = c("k", "estimator"))
df2 <- left_join(df2, nnp, by = c("k", "estimator"))

#Prepping dataset for ggplot
df2_20 <- df2 %>% filter(estimator != "fe" & k > 19)
df2_20$estimator <- as.factor(df2_20$estimator)
levels(df2_20$estimator) <- c(g = TeX("$\\widehat{\\gamma}$"), G = TeX("$\\widehat{\\Gamma}$"))

fig3 <- ggplot(data = df2_20) + geom_point(aes(x = k, y = median))
fig3 <- fig3 + geom_errorbar(aes(x = k, y = median, ymin = fp, ymax = nnp), width = 0) +
  facet_wrap(~ estimator, labeller = label_parsed)
fig3 <- fig3 + theme_bw() + theme(panel.grid.major = element_blank(), panel.grid.minor = element_blank()) +
  ylab("Error") + xlab("Periods (50 units per period)") + geom_hline(yintercept = 0)
pdf("fig3.pdf", width = 5, height = 3)
fig3
dev.off()

#Combine objects fig1, fig2, fig3 into a single image
pdf("three_little_figs.pdf", width = 5, height = 9)
grid::grid.draw(rbind(ggplotGrob(fig1), ggplotGrob(fig2), ggplotGrob(fig3), size = "last"))
dev.off()



# Some simple analysis of the estimates
midmiss.fe <- results_panel %>% group_by(k) %>% summarize(middlemiss = median(fe))
midmiss.g <- results_panel %>% group_by(k) %>% summarize(middlemiss = median(g))
midmiss.G <- results_panel %>% group_by(k) %>% summarize(middlemiss = median(G))
mysd.fe <- results_panel %>% group_by(k) %>% summarize(mysd = sd(fe))
mysd.g.panel <- results_panel %>% group_by(k) %>% summarize(mysd = sd(g))
mysd.G.panel <- results_panel %>% group_by(k) %>% summarize(mysd = sd(G))



#Code to create figure 3 in article
x <- c(32,28,42,38,30,34,26,40,46,34,36,34,26,24,46,44)
y <- c(28,32,38,42,30,26,34,40,34,46,34,36,24,26,44,46)
group <- c(rep("a", 4), rep("b", 6), rep("c", 6))
df <- tibble(x, y, group)
q <- ggplot(data = df) + geom_point(aes(x = x, y = y, shape = group, color = group), size = 2) +
  theme_classic()
q <- q + scale_color_manual(values = c("black", "black", "red")) + scale_shape_manual(values = c(19, 1, 19))
q <- q + xlab(TeX("$D_i$")) + ylab(TeX("$\\gamma_i$")) + theme(legend.position = "none") +
  theme(axis.title.y = element_text(angle = 0, size = 20)) +
  theme(axis.title.x = element_text(hjust = 1, size = 15)) +
  geom_abline(intercept = 0, slope = 1, linetype = 3) +
  xlim(22, 48) + ylim(22, 48) +
  theme(axis.text = element_blank()) +
  theme(axis.ticks = element_blank())
pdf(file = "little_gamma.pdf", width = 4, height = 4)
q
dev.off()

